import os
import json
import jieba
import torch
import requests
import numpy as np
from utils.log import MyLog
from pyltp import Postagger
#from transformers import AutoConfig, AutoModel, AutoTokenizer


class TitleInfo:
    def __init__(self, config):
        self.config = config
        self.log = MyLog(config.log_dir, __name__).getlog()
        self.postagger = Postagger()
        self.postagger.load(os.path.join(config.LTP_DIR, "pos.model"))
        self.textsim = torch.jit.load(config.LTP_DIR + 'textsim.pt')
        self.textsim.eval()
        self.COVID = config.COVID_RE
        self.stopwords = [i.strip() for i in
                          open(config.stopword_path + '/stop.txt', 'r', encoding='utf-8').readlines()]
        self.word2idx = json.loads(open(config.stopword_path + '/word2id.json', 'r', encoding='utf-8').read())

    def __del__(self):
        self.postagger.release()

    '''jieba分词去停用词'''

    def get_words(self, sentence):
        if not sentence: return []
        words = jieba.lcut(sentence)
        for i in words:
            if i in self.stopwords:
                words.remove(i)
            if i == ' ':
                words.remove(i)
        return list(filter(None, words))

    """获取标题分词结果,姓名,公司名,专有名词,英文标识"""

    def titleinfo(self, title):
        words = jieba.lcut(title)
        psg = list(self.postagger.postag(words))
        wp = {w: p for w, p in zip(words, psg)}
        words = list(wp.keys())
        names = [w for w, p in wp.items() if p.startswith('nh')]
        areas = [w for w, p in wp.items() if p.startswith('ns')]
        comp = [w for w, p in wp.items() if p.startswith('nz')]
        abb = [w for w, p in wp.items() if p.startswith('j')]
        eng = [w for w, p in wp.items() if p.startswith('ws')]
        return {'words': words, 'names': names, 'areas': areas, 'comp': comp, 'abb': abb, 'eng': eng}

    """抽取标题向量"""

    def title2vec(self, title):
        with torch.no_grad():
            encode = self.tokenizer.encode_plus(title,
                                                padding=True,
                                                truncation=True,
                                                max_length=100,
                                                return_tensors="pt")
            output = self.model(encode['input_ids'], encode['token_type_ids'], encode['attention_mask'])
            vec = output[0].squeeze(0).mean(0,keepdim=True).view(-1)
            return vec.numpy()

    # def title2vec(self, title):
    #     query = {"text": title}
    #     try:
    #         res = requests.get(self.config.vecurl, params=query)
    #         vec = json.loads(res.text).get('ret')
    #     except:
    #         pass
    #     return np.array(eval(vec))

    """使用Esim获取标题相似度"""

    def getsim(self, title1, title2):
        t1 = [self.word2idx.get(i, 1) for i in title1][:16]
        t2 = [self.word2idx.get(i, 1) for i in title2][:16]
        t1.extend([0] * (16 - len(t1)))
        t2.extend([0] * (16 - len(t2)))
        t1, t2 = torch.tensor([t1]), torch.tensor([t2])
        t3 = torch.stack((t1, t2), 0)
        pred = self.textsim(t3)
        sim = torch.log(pred[0][1] / pred[0][0]).item()
        return sim

    def rule_limit(self, title, r_title, info, r_info):
        sim = self.getsim(title, r_title)
        ori_covid, rec_covid = bool(self.COVID.search(title)), bool(self.COVID.search(r_title))
        if ori_covid or rec_covid:
            if ori_covid and rec_covid:
                if info.get('areas') and r_info.get('areas'):
                    if not set(info.get('areas')).intersection(set(r_info.get('areas'))):
                        return None
            else:
                sim *= 0.1
        if info.get('comp') != r_info.get('comp'):
            sim *= 0.15
        if info.get('names') or r_info.get('names'):
            if len(set(''.join(info.get('names')).strip()).intersection(set(''.join(r_info.get('names')).strip()))) > 1:
                sim = abs(sim) * 1.8
            elif info.get('names') and r_info.get('names'):
                return None
            else:
                sim *= 0.5

        if info.get('eng') and r_info.get('eng'):
            if set(info.get('eng')).intersection(set(r_info.get('eng'))):
                sim *= 1.3
            else:
                sim *= 0.8
        if info.get('areas') and r_info.get('areas'):
            if set(info.get('areas')).intersection(set(r_info.get('areas'))):
                sim *= 1.2
            else:
                return None
        return sim if sim > -1 else None


if __name__ == '__main__':
    pass

